/*
 *
 * Copyright (c) 2022 The Raptor Authors. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

#include "src/windows/tcp_client.h"

#include <string.h>

#include "raptor-lite/impl/connector.h"
#include "raptor-lite/impl/config.h"
#include "raptor-lite/impl/endpoint.h"
#include "raptor-lite/utils/log.h"
#include "raptor-lite/utils/time.h"
#include "src/common/endpoint_impl.h"
#include "src/common/socket_util.h"
#include "src/windows/socket_settings.h"

namespace raptor {

struct InternalClientEventNode {
    MultiProducerSingleConsumerQueue::Node node;
    NetEvent event;
    int err_code;
    Slice data;
};

#define close_socket(fd)                                                                           \
    if (fd != INVALID_SOCKET) {                                                                    \
        ::closesocket(fd);                                                                         \
        fd = INVALID_SOCKET;                                                                       \
    }

constexpr int TEMP_BUFFER_SIZE = 4 * 1024 * 1024;
constexpr int MIN_TIME_SLICE = 300;

TcpClient::TcpClient(ClientHandler *handler)
    : service_(handler)
    , shutdown_(true)
    , is_connected_(false)
    , send_pending_(false)
    , fd_(INVALID_SOCKET)
    , connectex_(NULL)
    , next_package_size_(NULL)
    , last_sendtime_(0)
    , last_recvtime_(0)
    , count_(0) {

    buffer_ = new char[TEMP_BUFFER_SIZE];
    conn_ole_ = new OverLappedEx;
    send_ole_ = new OverLappedEx;
    recv_ole_ = new OverLappedEx;
}

TcpClient::~TcpClient() {
    if (buffer_) delete[] buffer_;
    if (conn_ole_) delete conn_ole_;
    if (send_ole_) delete send_ole_;
    if (recv_ole_) delete recv_ole_;
}

raptor_error TcpClient::Init() {

    if (!shutdown_) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient is already running");
    }

    raptor_error e = iocp_.create(1);
    if (e != RAPTOR_ERROR_NONE) {
        return e;
    }

    bool success = false;
    mpscq_thd_ =
        Thread("client_thread", std::bind(&TcpClient::WorkThread, this, std::placeholders::_1),
               nullptr, &success);

    if (!success) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient failed to create mpscq thread");
    }

    iocp_thd_ =
        Thread("iocp_thread", std::bind(&TcpClient::IocpProcess, this, std::placeholders::_1),
               nullptr, &success);

    if (!success) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient failed to create iocp thread");
    }

    memset(conn_ole_, 0, sizeof(OverLappedEx));
    memset(send_ole_, 0, sizeof(OverLappedEx));
    memset(recv_ole_, 0, sizeof(OverLappedEx));

    conn_ole_->event_type = internal::kConnectEvent;
    send_ole_->event_type = internal::kSendEvent;
    recv_ole_->event_type = internal::kRecvEvent;

    shutdown_ = false;

    return RAPTOR_ERROR_NONE;
}

raptor_error TcpClient::Start() {
    if (shutdown_) {
        return RAPTOR_ERROR_FROM_STATIC_STRING("TcpClient is not initialized");
    }

    count_.Store(0);
    next_package_size_ = 0;
    last_sendtime_ = GetCurrentMilliseconds();
    last_recvtime_ = last_sendtime_;
    mpscq_thd_.Start();
    iocp_thd_.Start();
    return RAPTOR_ERROR_NONE;
}

void TcpClient::Shutdown() {
    if (!shutdown_) {
        log_warn("TcpClient: prepare to shutdown");
        shutdown_ = true;

        close_socket(fd_);
        is_connected_ = false;

        cv_.Signal();
        iocp_thd_.Join();
        mpscq_thd_.Join();
        iocp_.shutdown();

        // clear message queue
        bool empty = true;
        do {
            auto n = mpscq_.PopAndCheckEnd(&empty);
            if (n != nullptr) {
                delete reinterpret_cast<InternalClientEventNode *>(n);
            }
        } while (!empty);

        rcv_buffer_.ClearBuffer();
        snd_buffer_.ClearBuffer();
        log_warn("TcpClient: shutdown");
    }
}

raptor_error TcpClient::Connect(const std::string &addr) {
    raptor_resolved_addresses *addrs;
    auto e = raptor_blocking_resolve_address(addr.c_str(), nullptr, &addrs);
    if (e != RAPTOR_ERROR_NONE) {
        return e;
    }
    if (addrs->naddrs == 0) {
        return RAPTOR_ERROR_FROM_FORMAT("Invalid address: (%s)", addr.c_str());
    }

    for (int i = 0; i < static_cast<int>(addrs->naddrs); i++) {
        e = AsyncConnect(&addrs->addrs[i], 3000);
        if (e != RAPTOR_ERROR_NONE) {
            log_warn("TcpClient: Failed to connect %s, %s", addr.c_str(), e->ToString().c_str());
        } else {
            break;
        }
    }
    raptor_resolved_addresses_destroy(addrs);
    return e;
}

raptor_error TcpClient::AsyncConnect(const raptor_resolved_address *addr, int tcp_user_timeout_ms) {

    raptor_dualstack_mode mode = RAPTOR_DSMODE_NONE;
    raptor_resolved_address local_address;
    raptor_resolved_address mapped_addr;
    int status;
    BOOL ret;
    SOCKET fd;
    char *str_addr = nullptr;

    // SHADOW: If OVERLAPPED not initialized, ConnectEx will fail
    // and WSAGetLastError will return 6.
    memset(conn_ole_, 0, sizeof(*conn_ole_));
    conn_ole_->event_type = internal::kConnectEvent;

    raptor_error error = raptor_create_socket(addr, &mapped_addr, &fd, &mode);

    if (error != RAPTOR_ERROR_NONE) {
        goto failure;
    }
    error = raptor_tcp_prepare_socket(fd, tcp_user_timeout_ms);
    if (error != RAPTOR_ERROR_NONE) {
        goto failure;
    }
    error = GetConnectExIfNecessary(fd);
    if (error != RAPTOR_ERROR_NONE) {
        goto failure;
    }

    raptor_sockaddr_make_wildcard6(0, &local_address);

    status = bind(fd, (raptor_sockaddr *)&local_address.addr, (int)local_address.len);

    if (status != 0) {
        error = RAPTOR_WINDOWS_ERROR(WSAGetLastError(), "bind");
        goto failure;
    }

    iocp_.add(fd, this);

    fd_ = fd;

    ret = connectex_(fd, (raptor_sockaddr *)&mapped_addr.addr, (int)mapped_addr.len, NULL, 0, NULL,
                     &conn_ole_->overlapped);

    if (ret) {
        if (!AsyncRecv()) {
            error = RAPTOR_WINDOWS_ERROR(WSAGetLastError(), "WSARecv");
            goto failure;
        }
        last_sendtime_ = GetCurrentMilliseconds();
        last_recvtime_ = last_sendtime_;
        OnConnected();

    } else {
        /* It wouldn't be unusual to get a success immediately. But we'll still get
         an IOCP notification, so let's ignore it. */
        int last_error = WSAGetLastError();
        if (last_error != ERROR_IO_PENDING) {
            error = RAPTOR_WINDOWS_ERROR(WSAGetLastError(), "ConnectEx");
            goto failure;
        }
    }

    raptor_sockaddr_to_string(&str_addr, &mapped_addr, 0);
    if (str_addr) {
        log_info("TcpClient: start connecting %s", str_addr);
        free(str_addr);
    }
    return RAPTOR_ERROR_NONE;

failure:
    if (fd != INVALID_SOCKET) {
        closesocket(fd);
    }
    fd_ = INVALID_SOCKET;
    return error;
}

void TcpClient::IocpProcess(void *) {
    while (!shutdown_) {
        DWORD NumberOfBytesTransferred = 0;
        void *CompletionKey = NULL;
        LPOVERLAPPED lpOverlapped = NULL;

        // https://docs.microsoft.com/en-us/windows/win32/api/ioapiset/nf-ioapiset-getqueuedcompletionstatus
        bool ret = iocp_.polling(&NumberOfBytesTransferred, (PULONG_PTR)&CompletionKey,
                                 &lpOverlapped, MIN_TIME_SLICE);

        if (!ret) {
            /*
                If theGetQueuedCompletionStatus function succeeds, it dequeued a completion
                packet for a successful I/O operation from the completion port and has
                stored information in the variables pointed to by the following parameters:
                lpNumberOfBytes, lpCompletionKey, and lpOverlapped. Upon failure (the return
                value is FALSE), those same parameters can contain particular value
                combinations as follows:

                    (1) If *lpOverlapped is NULL, the function did not dequeue a completion
                        packet from the completion port. In this case, the function does not
                        store information in the variables pointed to by the lpNumberOfBytes
                        and lpCompletionKey parameters, and their values are indeterminate.

                    (2) If *lpOverlapped is not NULL and the function dequeues a completion
                        packet for a failed I/O operation from the completion port, the
                        function stores information about the failed operation in the
                        variables pointed to by lpNumberOfBytes, lpCompletionKey, and
                        lpOverlapped. To get extended error information, call GetLastError.
            */

            if (lpOverlapped != NULL && CompletionKey != NULL) {
                // Maybe an error occurred or the connection was closed

                OverLappedEx *olex = (OverLappedEx *)lpOverlapped;

                int err = WSAGetLastError();

                OnErrorEvent(olex->event_type, err);

            } else {
                // IOCP TIMEOUT
            }
            continue;
        }

        OverLappedEx *olex = (OverLappedEx *)lpOverlapped;

        if (!is_connected_ && olex->event_type & internal::kConnectEvent) {
            if (!AsyncRecv()) {
                OnErrorEvent(internal::kRecvEvent, WSAGetLastError());
            } else {
                last_sendtime_ = GetCurrentMilliseconds();
                last_recvtime_ = last_sendtime_;
                OnConnected();
            }
            continue;
        }

        if (olex->event_type & internal::kSendEvent) {
            this->OnSendEvent(NumberOfBytesTransferred);
        }
        if (olex->event_type & internal::kRecvEvent) {
            this->OnRecvEvent(NumberOfBytesTransferred);
        }
    }
}

void TcpClient::OnErrorEvent(int event_type, int err) {
    log_info("TcpClient error event, event_type:%d err:%d", event_type, err);
    close_socket(fd_);
    auto obj = new InternalClientEventNode;
    if (is_connected_) {
        is_connected_ = false;
        obj->event = NetEvent::Closed;
    } else {
        obj->event = NetEvent::ConnectFailed;
    }
    obj->err_code = err;
    mpscq_.push(&obj->node);
    count_.FetchAdd(1, MemoryOrder::RELAXED);
    cv_.Signal();
}

void TcpClient::OnConnected() {
    is_connected_ = true;
    auto obj = new InternalClientEventNode;
    obj->event = NetEvent::ConnectSuccess;
    obj->err_code = 0;
    mpscq_.push(&obj->node);
    count_.FetchAdd(1, MemoryOrder::RELAXED);
    cv_.Signal();
}

void TcpClient::OnSendEvent(DWORD send_bytes) {
    if (send_bytes == 0) return;
    if (!SendImpl(send_bytes)) {
        int err = WSAGetLastError();
        OnErrorEvent(internal::kSendEvent, err);
        return;
    }
}

void TcpClient::OnRecvEvent(DWORD recv_bytes) {
    if (!RecvImpl(recv_bytes) || !AsyncRecv()) {
        int err = WSAGetLastError();
        OnErrorEvent(internal::kRecvEvent, err);
        return;
    }
}

bool TcpClient::SendMsg(const Slice &s) {
    if (shutdown_ || fd_ == INVALID_SOCKET || s.empty()) {
        return false;
    }
    AutoMutex g(&smtx_);
    snd_buffer_.AddSlice(s);
    return AsyncSend();
}

void TcpClient::Disconnect() { close_socket(fd_); }

bool TcpClient::SendImpl(size_t size) {
    AutoMutex g(&smtx_);
    send_pending_ = false;
    snd_buffer_.MoveHeader(size);
    last_sendtime_ = GetCurrentMilliseconds();
    return AsyncSend();
}

constexpr size_t MAX_PACKAGE_SIZE = 4 * 1024 * 1024;
constexpr size_t MAX_WSABUF_COUNT = 24;

bool TcpClient::AsyncSend() {
    if (snd_buffer_.Empty() || send_pending_) {
        return true;
    }

    WSABUF buffers[MAX_WSABUF_COUNT];
    size_t prepare_send_size = 0;

    send_pending_ = true;

    size_t i, count = snd_buffer_.SliceCount();
    for (i = 0; i < count && i < MAX_WSABUF_COUNT; i++) {
        const Slice &s = snd_buffer_[i];

        if (i != 0 && prepare_send_size + s.length() > MAX_PACKAGE_SIZE) {
            break;
        }

        buffers[i].buf = s.buffer();
        buffers[i].len = static_cast<ULONG>(s.length());
        prepare_send_size += s.length();
    }

    // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-wsasend
    int ret = WSASend(fd_, buffers, (DWORD)i, NULL, 0, &send_ole_->overlapped, NULL);

    if (ret == SOCKET_ERROR) {
        int error = WSAGetLastError();
        if (error != WSA_IO_PENDING) {
            send_pending_ = false;
            return false;
        }
    }
    return true;
}

bool TcpClient::AsyncRecv() {
    DWORD dwFlags = 0;
    WSABUF buffer;

    buffer.buf = NULL;
    buffer.len = 0;

    // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-wsarecv
    int ret = WSARecv(fd_, &buffer, 1, NULL, &dwFlags, &recv_ole_->overlapped, NULL);

    if (ret == SOCKET_ERROR) {
        int error = WSAGetLastError();
        if (error != WSA_IO_PENDING) {
            return false;
        }
    }
    return true;
}

bool TcpClient::RecvImpl(size_t size) {
    (size);

    // AutoMutex g(&_rcv_mtx);
    int unused_space = TEMP_BUFFER_SIZE;
    int recv_bytes = 0;

    do {

        recv_bytes = ::recv(fd_, buffer_, unused_space, 0);

        if (recv_bytes == 0) return false;
        if (recv_bytes < 0) {
            int err = WSAGetLastError();
            if (WSAEWOULDBLOCK == err) {
                break;
            }
            return false;
        }

        last_recvtime_ = GetCurrentMilliseconds();

        TransferData(Slice(buffer_, recv_bytes));

    } while (recv_bytes == unused_space);

    return true;
}

void TcpClient::TransferData(const Slice &s) {
    auto obj = new InternalClientEventNode;
    obj->event = NetEvent::Null;
    obj->err_code = 0;
    obj->data = s;
    mpscq_.push(&obj->node);
    count_.FetchAdd(1, MemoryOrder::RELAXED);
    cv_.Signal();
}

void TcpClient::WorkThread(void *) {
    while (!shutdown_) {
        if (count_.Load() < 1) {
            AutoMutex g(&cmtx_);
            cv_.Wait(&cmtx_, 1000);
            continue;
        }
        auto n = mpscq_.pop();
        if (n != nullptr) {
            count_.FetchSub(1, MemoryOrder::RELAXED);
            auto msg = reinterpret_cast<InternalClientEventNode *>(n);
            Dispatch(msg);
            delete msg;
        }
    }
}

void TcpClient::Dispatch(InternalClientEventNode *msg) {
    if (!msg->data.empty()) {
        rcv_buffer_.AddSlice(msg->data);
        if (rcv_buffer_.GetBufferLength() >= next_package_size_) {
            ParsingProtocol();
        }
        return;
    }

    service_->OnNetEvent(msg->event);
}

Slice TcpClient::ReadSliceFromRecvBuffer(size_t c) {
    Slice ret;
    size_t count = rcv_buffer_.SliceCount();
    for (size_t i = 0; i < count && i < c; i++) {
        ret += rcv_buffer_[i];
    }
    return ret;
}

void TcpClient::ParsingProtocol() {

    size_t cache_size = rcv_buffer_.GetBufferLength();
    int package_counter = 0;

    if (cache_size >= next_package_size_ && next_package_size_ > 0) {
        Slice packet = rcv_buffer_.GetHeader(next_package_size_);
        service_->OnMessage(packet);
        rcv_buffer_.MoveHeader(next_package_size_);
        cache_size -= next_package_size_;
        next_package_size_ = 0;
    }

    while (cache_size > 0) {
        size_t read_count = 1;
        int pack_len = 0;
        Slice packet;
        do {
            packet = ReadSliceFromRecvBuffer(read_count);
            pack_len = service_->MessageLength((uint8_t *)packet.buffer(), packet.length());
            if (pack_len < 0) {
                // discard buffer
                rcv_buffer_.MoveHeader(packet.length());
                cache_size = rcv_buffer_.GetBufferLength();
                break;
            }

            // equal 0 means we need more data
            if (pack_len == 0) {
                // reach the end of rcv_buffer
                if (read_count >= rcv_buffer_.SliceCount()) {
                    goto done;
                }
                // read more data to check msg header
                read_count++;
                continue;
            }

            // set next_package_size_ for next time
            if (cache_size < static_cast<size_t>(pack_len)) {
                next_package_size_ = pack_len;
                goto done;
            }

            // We got a whole packet with length pack_len
            break;

        } while (true);

        if (pack_len < 0) continue;

        if (packet.size() < static_cast<size_t>(pack_len)) {
            packet = rcv_buffer_.GetHeader(pack_len);
        } else {
            size_t n = packet.size() - pack_len;
            packet.RemoveTail(n);
        }
        service_->OnMessage(packet);
        rcv_buffer_.MoveHeader(pack_len);

        cache_size = rcv_buffer_.GetBufferLength();
    }
done:
    (void)package_counter;
}

raptor_error TcpClient::GetConnectExIfNecessary(SOCKET s) {
    if (!connectex_) {
        GUID guid = WSAID_CONNECTEX;
        DWORD ioctl_num_bytes;
        int status;

        /* Grab the function pointer for ConnectEx for that specific socket.
            It may change depending on the interface. */
        status = WSAIoctl(s, SIO_GET_EXTENSION_FUNCTION_POINTER, &guid, sizeof(guid), &connectex_,
                          sizeof(connectex_), &ioctl_num_bytes, NULL, NULL);

        if (status != 0) {
            return RAPTOR_WINDOWS_ERROR(WSAGetLastError(),
                                        "WSAIoctl(SIO_GET_EXTENSION_FUNCTION_POINTER)");
        }
    }

    return RAPTOR_ERROR_NONE;
}
} // namespace raptor
